package fem2.strategies;

import fem2.Debugger;
import fem2.Model;
import fem2.enu.EchoLevelType;
import fem2.enu.MassMatrixType;
import fem2.enu.NewmarkType;
import fem2.jlinalg.DiagonalMatrix;
import fem2.jlinalg.Solver;
import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;

public abstract class NewmarkCrisfieldFormulationTimeIntegrationStrategy extends
		NewmarkTimeIntegrationStrategy {

	private ArrayVector ri;

	public NewmarkCrisfieldFormulationTimeIntegrationStrategy(Model m, MassMatrixType massType,
			double alpha1, double alpha2, double alphaf, double alpham, double beta, double gamma) {
		super(m, massType, alpha1, alpha2, alphaf, alpham, beta, gamma);
	}

	public NewmarkCrisfieldFormulationTimeIntegrationStrategy(Model m, MassMatrixType massType,
			double alpha1, double alpha2, NewmarkType type, double rhoInf) {
		super(m, massType, alpha1, alpha2, type, rhoInf);
	}

	@Override
	public void started(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		super.started(leftHandSideSolver, rightHandSide, u);

		int n = model.getSize();
		this.ri = new ArrayVector(n);
	}

	@Override
	public void timeStepStarted(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		super.timeStepStarted(leftHandSideSolver, rightHandSide, u);

		/*
		 * compute first ri for the time step. Actually it's the ri for the last
		 * converged time step.
		 */
		model.assembleRi(this.ri.getValues());

		if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
			Debugger.watch("ri = ", ri);
		}
	}

	@Override
	public void assembleEffectiveSystem(IMatrix kteff, IVector reff) {
		int n = model.getSize();
		double dt = t - t_old;

		// Debugger.watch("uHat = ", uHat);
		// Debugger.watch("uHat_old = ", uHat_old);
		// Debugger.watch("udHat_old = ", udHat_old);
		// Debugger.watch("uddHat_old = ", uddHat_old);

		/*
		 * set the uHatGlobal for model. For Crisfield formulation we don't use
		 * u_n_plus_1_minus_alphaf to apply in the system
		 */
		// model.setX(t, uHat.getValues());

		/*
		 * compute M and Kt(u)
		 */
		BLAM.zero(kteff);

		IMatrix M;
		if (massType == MassMatrixType.DIRECT_MASS_LUMPING) {
			M = new DiagonalMatrix(n);
		} else {
			M = (IMatrix) kteff.clone();
		}

		model.assembleKt(kteff);
		model.assembleM(M, massType);

		// Debugger.watch("alpham = ", alpham);
		// Debugger.watch("alphaf = ", alphaf);
		// Debugger.watch("beta = ", beta);
		// Debugger.watch("gamma = ", gamma);
		// Debugger.watch("kt = ", kteff);
		// Debugger.watch("M = ", M);

		/****************************************
		 * compute reff. At this point keff play the role of kt
		 ****************************************/
		/*
		 * 1: compute r_(n+1-alphaf)
		 */
		model.assembleR(1 - alphaf, ((ArrayVector) reff).getValues(), t);
		model.assembleR(alphaf, ((ArrayVector) reff).getValues(), t_old);

		/*
		 * 2: compute reff = reff - (1-alphaf)ri(u) - alphaf ri(un) TODO: check
		 * if we can optimize by introducing assembleRi(alpha,ri)
		 */
		ArrayVector ri = new ArrayVector(n);
		model.assembleRi(ri.getValues());
		BLAM.add(-(1.0 - alphaf), ri, reff);
		BLAM.add(-alphaf, this.ri, reff);

		/*
		 * 3: compute Reff = Reff - D*ud_(n+1-alphaf)(ud_(n+1)(u_(n+1))). At
		 * this point u_(n+1) = uHat
		 */
		IVector ud_n_plus_1_minus_alphaf = (IVector) uHat.clone();
		BLAM.add(-1.0, uHat_old, ud_n_plus_1_minus_alphaf);
		double coeff = ((1 - alphaf) * gamma / beta) / dt;
		BLAM.scale(coeff, ud_n_plus_1_minus_alphaf);
		BLAM.add(1.0 - coeff * dt, udHat_old, ud_n_plus_1_minus_alphaf);
		BLAM.add(-(1 - alphaf) * (gamma - 2 * beta) * dt / (2 * beta), uddHat_old,
				ud_n_plus_1_minus_alphaf);
		BLAM.multiply(-alpha1, BLAM.NO_TRANSPOSE, M, ud_n_plus_1_minus_alphaf, 1.0, reff);
		BLAM.multiply(-alpha2, BLAM.NO_TRANSPOSE, kteff, ud_n_plus_1_minus_alphaf, 1.0, reff);
		/*
		 * 4: compute Reff = Reff - M*udd_(n+1-alpham)
		 */
		IVector udd_n_plus_1_minus_alpham = (IVector) uHat.clone();
		BLAM.add(-1.0, uHat_old, udd_n_plus_1_minus_alpham);
		coeff = ((1 - alpham) / beta) / Math.pow(dt, 2);
		BLAM.scale(coeff, udd_n_plus_1_minus_alpham);
		BLAM.add(-coeff * dt, udHat_old, udd_n_plus_1_minus_alpham);
		BLAM.add(1.0 - (1 - alpham) / (2 * beta), uddHat_old, udd_n_plus_1_minus_alpham);
		BLAM.multiply(-1.0, BLAM.NO_TRANSPOSE, M, udd_n_plus_1_minus_alpham, 1.0, reff);

		/*
		 * compute Keff
		 */
		double dampingCoeff = (gamma * (1 - alphaf) / beta) / dt;
		double massCoeff = ((1 - alpham) / beta) / Math.pow(dt, 2);
		BLAM.scale(1 - alphaf + alpha2 * dampingCoeff, kteff);
		BLAM.add(massCoeff + alpha1 * dampingCoeff, M, kteff);

		if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
			Debugger.watch("kteff = ", kteff);
			Debugger.watch("reff = ", reff);
		}
	}
}
